package org.kiwi.data.config;

import org.apache.commons.lang3.ArrayUtils;
import org.kiwi.data.RoutingDataSource;
import org.kiwi.data.enums.LbType;
import org.kiwi.data.interceptor.ReadWriteSeparationInterceptor;
import org.kiwi.data.lb.RandomLoadBalanceStrategy;
import org.kiwi.data.lb.RoundRobinLoadBalanceStrategy;
import org.kiwi.data.util.AopUtil;
import org.kiwi.data.util.BeanDefinitionParserUtil;
import org.kiwi.data.vo.DataSourceMetaVO;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.RuntimeBeanReference;
import org.springframework.beans.factory.support.ManagedMap;
import org.springframework.beans.factory.support.RootBeanDefinition;
import org.springframework.beans.factory.xml.BeanDefinitionParser;
import org.springframework.beans.factory.xml.ParserContext;
import org.w3c.dom.Element;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.Map;

import static org.kiwi.data.cons.DataConstant.*;

/**
 * Created by jack08.liu on 2016/12/21.
 */
public class ReadWriteSeparationBeanDefinitionParser implements BeanDefinitionParser {

    private Map<String, Map<String, Integer>> groupLbFactorsMap = new LinkedHashMap<>();
    private Map<String, String> writeDsBeanNameMap = new LinkedHashMap<>();

    private ManagedMap<String, RuntimeBeanReference> loadBalanceStrategyMap = new ManagedMap<>();

    @Override
    public BeanDefinition parse(Element element, ParserContext parserContext) {
        try {
            DataSourceMetaVO dataSourceMetaVO = BeanDefinitionParserUtil.resolve(element, parserContext);

            setUpMetaInfo(parserContext, dataSourceMetaVO);

            registerLoadBalancer(parserContext, dataSourceMetaVO);

            registerReadWriteSeparationInterceptor(parserContext);

            registerReadWriteSeparationAopConfig(element, parserContext, element.getAttribute(ATTRIBUTE_POINTCUT_EXPRESSION));

            updateRoutingDataSource(parserContext);

            return null;
        } catch (Exception e) {
            parserContext.getReaderContext().error(
                    "register ReadWriteSeparationBeanDefinition has some problem:[" + e.getMessage() + "]",
                    parserContext.extractSource(element));
            throw e;
        }
    }

    private void setUpMetaInfo(ParserContext parserContext, DataSourceMetaVO dataSourceMetaVO) {
        BeanDefinitionParserUtil.setUpAtomicDataSourceMetaInfo(
                parserContext,
                dataSourceMetaVO,
                this.writeDsBeanNameMap,
                this.groupLbFactorsMap
        );
    }

    private void registerLoadBalancer(ParserContext parserContext, DataSourceMetaVO dataSourceMetaVO) {
        for (DataSourceMetaVO.GroupDataSourceMetaVO groupDataSourceMetaVO : dataSourceMetaVO.getGroups()) {
            doRegisterLoadBalancer(parserContext, groupDataSourceMetaVO);
            String groupName = groupDataSourceMetaVO.getGroupName();
            this.loadBalanceStrategyMap.put(groupName, new RuntimeBeanReference(groupName + "_" + LOAD_BALANCER));
        }
    }

    private void doRegisterLoadBalancer(ParserContext parserContext, DataSourceMetaVO.GroupDataSourceMetaVO groupDataSourceMetaVO) {
        String groupName = groupDataSourceMetaVO.getGroupName();

        RootBeanDefinition loadBalancer = new RootBeanDefinition();
        LbType lbType = LbType.determineLbTypeByText(groupDataSourceMetaVO.getLoadBalance());
        if (lbType == LbType.ROUNDROBIN) {
            loadBalancer.setBeanClass(RoundRobinLoadBalanceStrategy.class);
            loadBalancer.getConstructorArgumentValues().addIndexedArgumentValue(0,
                    this.groupLbFactorsMap.get(groupName));
        } else if (lbType == LbType.RANDOM) {
            loadBalancer.setBeanClass(RandomLoadBalanceStrategy.class);
            loadBalancer.getConstructorArgumentValues().addIndexedArgumentValue(0,
                    new ArrayList<>(this.groupLbFactorsMap.get(groupName).keySet()));
        } else {
            throw new UnsupportedOperationException("unsupported lb type");
        }
        loadBalancer.getConstructorArgumentValues().addIndexedArgumentValue(1, this.writeDsBeanNameMap.get(groupName));
        parserContext.getRegistry().registerBeanDefinition(groupName + "_" + LOAD_BALANCER, loadBalancer);
    }

    private void registerReadWriteSeparationInterceptor(ParserContext parserContext) {
        RootBeanDefinition readWriteSeparationInterceptor = new RootBeanDefinition(ReadWriteSeparationInterceptor.class);
        readWriteSeparationInterceptor.getPropertyValues().add(LOAD_BALANCE_STRATEGY_MAP, this.loadBalanceStrategyMap);
        parserContext.getRegistry().registerBeanDefinition(READ_WRITE_SEPARATION_INTERCEPTOR, readWriteSeparationInterceptor);
    }

    private void registerReadWriteSeparationAopConfig(Element element, ParserContext parserContext, String pointCutExpression) {
        AopUtil.registerAopConfig(
                element,
                parserContext,
                READ_WRITE_SEPARATION_POINTCUT,
                pointCutExpression,
                READ_WRITE_SEPARATION_INTERCEPTOR,
                READ_WRITE_SEPARATION_ADVICE_ORDER
        );
    }

    private void updateRoutingDataSource(ParserContext parserContext) {
        String[] beanNames = parserContext.getRegistry().getBeanDefinitionNames();
        if (ArrayUtils.isNotEmpty(beanNames)) {
            for (String beanName : beanNames) {
                BeanDefinition beanDefinition = parserContext.getRegistry().getBeanDefinition(beanName);
                if (beanDefinition instanceof RootBeanDefinition) {
                    RootBeanDefinition rootBeanDefinition = (RootBeanDefinition) beanDefinition;
                    try {
                        if (rootBeanDefinition.getBeanClass() != null && RoutingDataSource.class.isAssignableFrom(rootBeanDefinition.getBeanClass())) {
                            rootBeanDefinition.getPropertyValues().add(READ_WRITE_ENABLED, true);
                        }
                    } catch (Exception e) {
                        //ignore
                        continue;
                    }
                }
            }
        }
    }

}
